from django.db import migrations


def migrate(apps, schema_editor):
    IngestorConfig = apps.get_model("ingestors_manager", "IngestorConfig")
    PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask")

    ic_threatfox = IngestorConfig.objects.get(name="ThreatFox")
    ic_malwarebazaar = IngestorConfig.objects.get(name="MalwareBazaar")
    pt_threatfox = PeriodicTask.objects.get(name="ThreatFoxIngestor")
    pt_malwarebazaar = PeriodicTask.objects.get(name="MalwareBazaarIngestor")

    ic_threatfox.disabled = True
    ic_malwarebazaar.disabled = True
    pt_threatfox.enabled = False
    pt_malwarebazaar.enabled = False

    ic_threatfox.routing_key = "ingestor"
    ic_malwarebazaar.routing_key = "ingestor"

    pt_threatfox.queue = "ingestor"
    pt_malwarebazaar.queue = "ingestor"

    ic_threatfox.full_clean()
    ic_malwarebazaar.full_clean()
    pt_threatfox.full_clean()
    pt_malwarebazaar.full_clean()

    ic_threatfox.save()
    ic_malwarebazaar.save()
    pt_threatfox.save()
    pt_malwarebazaar.save()


def reverse_migrate(apps, schema_editor):
    IngestorConfig = apps.get_model("ingestors_manager", "IngestorConfig")
    PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask")

    ic_threatfox = IngestorConfig.objects.get(name="ThreatFox")
    ic_malwarebazaar = IngestorConfig.objects.get(name="MalwareBazaar")
    pt_threatfox = PeriodicTask.objects.get(name="ThreatFoxIngestor")
    pt_malwarebazaar = PeriodicTask.objects.get(name="MalwareBazaarIngestor")

    ic_threatfox.routing_key = "default"
    ic_malwarebazaar.routing_key = "default"

    pt_threatfox.queue = "default"
    pt_malwarebazaar.queue = "default"

    ic_threatfox.full_clean()
    ic_malwarebazaar.full_clean()
    pt_threatfox.full_clean()
    pt_malwarebazaar.full_clean()

    ic_threatfox.save()
    ic_malwarebazaar.save()
    pt_threatfox.save()
    pt_malwarebazaar.save()


class Migration(migrations.Migration):
    atomic = False
    dependencies = [
        ("api_app", "0062_alter_parameter_python_module"),
        ("ingestors_manager", "0020_ingestor_config_threatfox"),
    ]

    operations = [migrations.RunPython(migrate, reverse_migrate)]
